source /root/miniconda3/etc/profile.d/conda.sh
conda activate ws
atk_index=$1
case_name=fused_infer_attention_score
rm -rf /etc/Ascend/ascend_cann_install.info

mkdir /home/ws/${atk_index}
mkdir /home/ws/${atk_index}/cann-package
mkdir /home/ws/${atk_index}/cann-ops-adv-dev
mkdir /home/ws/${atk_index}/result

# 安装cann
install_path=/home/ws/${atk_index}/cann-package
chmod -R 755 ${install_path}
export TMPDIR=/dev/shm
echo ">>>>>>>>>> start install cann <<<<<<<<<<"
/home/ws/Ascend-cann-toolkit_*_linux-aarch64.run --full --quiet --install-path=$install_path
/home/ws/Ascend-cann-kernels-*_linux-aarch64.run --install --quiet --install-path=$install_path
echo ">>>>>>>>>> end install cann <<<<<<<<<<"

# 编译
echo ">>>>>>>>>> start compile <<<<<<<<<<"
source /home/ws/${atk_index}/cann-package/ascend-toolkit/set_env.sh

cd /home/ws/${atk_index}
# cp -r /home/ws/cann-ops-adv-dev /home/ws/${atk_index}/

cd /home/ws/${atk_index}/cann-ops-adv-dev
bash build.sh -n incre_flash_attention -n prompt_flash_attention --disable-check-compatible
/home/ws/${atk_index}/cann-ops-adv-dev/output/CANN-custom_ops-*linux.aarch64.run --quiet

# 执行
echo ">>>>>>>>>> start run atk <<<<<<<<<<"
# cd /home/ws/${atk_index}
# git clone https://gitee.com/wangsong96/TempTools.git
# cp -r /home/ws/TempTools /home/ws/${atk_index}/
# cd /home/ws/${atk_index}/TempTools/DT/FusedIncreAttentionScoreV2
cd /home/ws/${atk_index}/cann-ops-adv-dev/tests/DT/${case_name}
source /home/ws/${atk_index}/cann-package/ascend-toolkit/set_env.sh
atk task -c ${case_name}.json -n node.yaml -p ${case_name}.py
cp -r /home/ws/${atk_index}/cann-ops-adv-dev/tests/DT/${case_name}/atk_output/${case_name}_*/* /home/ws/${atk_index}/result
# 判断是否通过atk
has_pass=$(grep -i -m1 -E "Pass" /home/ws/${atk_index}/result/log/atk.log | wc -l)
# [ $has_pass -gt 0 ] && echo true || echo false
if [ $has_pass -gt 0 ];then
    echo "ATK-test pass"
else
    echo "ATK-test faild"
    # tar -czvf /home/ws/${atk_index}/result_${build_index}.tar.gz /home/ws/${atk_index}/result
fi